# 蒙特卡洛树模块
import copy
import cmath
import math

import chess
from src.chess_util import Chess
import src.global_config as gl
import src.util as util
import src.evaluate as evaluate
import numpy as np
import multiprocessing as mp

chess_util = Chess()


# 节点类
class Node:
    def __init__(self, fen, step: str or None, parent, value=None):
        self.fen = fen
        # 最后回溯引擎推荐招法的时候可以用
        self.board = chess.Board(fen)
        self.step = step
        # 访问次数
        self.times = 0

        # 子节点列表
        self.leaf = []
        # 父节点
        self.parent = parent

        # 每次把估值算好
        if value is None:
            internal = chess_util.parse_internal_board_info(self.board, fen)
            self.value = chess_util.do_evaluate(self.board, internal)
            res = chess_util.check_board_status(self.board)
            self.is_end = res['end']
        else:
            self.value = value
            self.is_end = False


# 蒙特卡洛树搜索类
class Mcts:
    root_board = None
    board_list = None
    fen = None
    move = 0
    depth = 0

    # 获取蒙特卡洛树搜索相关的配置
    def __init__(self):
        self.mct_config = gl.get_value('mct_config')
        self.chess_util = chess_util
        self.get_config()

    # 获取配置文件
    def get_config(self):
        if self.mct_config is None:
            self.move = 2000
            self.depth = 15
        else:
            move = util.get_config_by_key(self.mct_config, 'move')
            depth = util.get_config_by_key(self.mct_config, 'depth')
            # 必须转int
            self.move = int(move)
            self.depth = int(depth)

    # 赋值 不想用初始化函数
    def set_value(self, board: chess.Board, fen: str):
        self.root_board: chess.Board = board
        self.board_list = self.chess_util.parse_board(board)
        self.fen = fen

    # 蒙特卡洛树搜索 分析
    def mct_search_analyze(self, fen: str):
        # 先创建初始局面
        root_board = chess.Board(fen)
        board = self.chess_util.parse_board(root_board)
        # 判断一进来就是棋局结束的特殊情况
        check = self.chess_util.check_board_status(root_board)
        if check['end'] is True:
            return None

        # 根节点
        root_node = Node(fen, None, None)
        color = self.chess_util.get_current_color_by_fen(fen)

        # 开始蒙特卡洛树搜索
        # 这里必须顺序执行 不能多线程! 也不一定
        for i in range(self.move):
            print("模拟[{}]次".format(i + 1))

            # 选择 如果已经到棋局尾声就跳出
            self.select(root_node, i, color)
        return self.get_best(root_node)

    # 处理完毕 获取最优节点
    def get_best(self, node: Node):
        leaf = node.leaf
        color = chess_util.get_current_color_by_fen(node.fen)
        if len(leaf) == 0:
            return None
        max_ucb = 0
        max_index = 0
        for i in range(len(leaf)):
            single = leaf[i]
            ucb = self.calc_UCB(single, self.move, color)
            # 白取值最大的 对黑来说最大值反而是ucb最小的
            if color is evaluate.Color.WHITE:
                max_ucb = max_ucb if max_ucb > ucb else ucb
                max_index = max_index if max_ucb > ucb else i
            else:
                max_ucb = max_ucb if max_ucb < ucb else ucb
                max_index = max_index if max_ucb < ucb else i
            print("ucb={},step={}".format(ucb, node.board.san(single.step)))

        choose = leaf[max_index]
        return choose.step, choose.fen

    # 选择
    def select(self, node: Node, num: int, color: evaluate.Color):
        # 不是叶子节点 找到当前节点最大的叶子节点进行选择
        if self.is_leaf(node) is False:
            if color is evaluate.Color.WHITE:
                max_ucb = 0
                max_index = 0
                for i in range(len(node.leaf)):
                    # 已经到棋局尾声的节点就跳过
                    if node.leaf[i].is_end is True:
                        continue
                    # 求最大值
                    leaf_ucb = self.calc_UCB(node.leaf[i], num, color)
                    # 无穷大就是最大
                    if math.isinf(leaf_ucb):
                        max_index = i
                        break
                    max_ucb = max_ucb if max_ucb > leaf_ucb else leaf_ucb
                    max_index = max_index if max_ucb > leaf_ucb else max_index
                return self.select(node.leaf[max_index], num, color)

            else:
                max_ucb = 0
                max_index = 0
                for i in range(len(node.leaf)):
                    # 求最大值
                    leaf_ucb = self.calc_UCB(node.leaf[i], num, color)
                    # 无穷大就是最大
                    if math.isinf(leaf_ucb):
                        max_index = i
                        break
                    max_ucb = max_ucb if max_ucb < leaf_ucb else leaf_ucb
                    max_index = max_index if max_ucb < leaf_ucb else max_index
                return self.select(node.leaf[max_index], num, color)

        # 叶子节点 进行仿真模拟
        else:
            node, value = self.extend(node, num)
            self.backward(node, num, value)

    # 扩展 根据深度进行扩展
    def extend(self, node: Node, num: int) -> tuple:
        board = node.board
        legal_list = list(board.legal_moves)
        value = 0
        legal_nums = len(legal_list)

        # 如果已经到棋局终点则不进行扩展
        if legal_nums == 0:
            value = self.chess_util.get_end_value(self.chess_util.check_board_status(board))
            node.is_end = True
            return node, value

        # 对每个节点按深度模拟一次
        for i in range(legal_nums):
            move = legal_list[i]
            temp_board = copy.deepcopy(board)
            temp_board.push(move)
            internal = chess_util.parse_internal_board_info(temp_board, temp_board.fen())
            get_value = chess_util.do_evaluate(temp_board, internal)
            if chess_util.get_current_color_by_fen(temp_board.fen()) is evaluate.Color.WHITE:
                value = value if value > get_value else value
            else:
                value = value if value < get_value else value
            node.leaf.append(Node(temp_board.fen(), move, node, value))
        return node, value

    #
    # # 多进程下的扩展逻辑
    # def extend_in_multiple(self,board,move):
    #     temp_board = copy.deepcopy(board)
    #     temp_board.push(move)
    #     internal = chess_util.parse_internal_board_info(temp_board, temp_board.fen())
    #     get_value = chess_util.do_evaluate(temp_board, internal)
    #     return get_value

    # 对单个节点进行模拟 得出结果
    def simulate(self, board: chess.Board, move, color: evaluate.Color):
        temp_board = copy.deepcopy(board)
        result = []
        for i in range(self.depth):

            legal_list = list(temp_board.legal_moves)
            legal_num = len(legal_list)
            values = []
            # 挨个模拟 放入数组
            # 如果棋局结束 就不能模拟了
            if legal_num == 0:
                value = chess_util.get_end_value(temp_board.result())
                values.append({
                    'value': value
                })

            for j in range(legal_num):
                work_board = copy.deepcopy(temp_board)
                work_board.push(legal_list[j])
                internal = chess_util.parse_internal_board_info(work_board, work_board.fen())
                value = chess_util.do_evaluate(work_board, internal)
                values.append({
                    'step': legal_list[j],
                    'value': value,
                    'board': work_board,
                })

            # 按颜色排序
            values = util.sort_list_by_key(values, 'value', color)
            best = values[0]
            result.append(best['value'])
            # 没必要继续模拟
            if legal_num == 0:
                break
            temp_board.push(best['step'])

        return np.mean(result)

    # 按深度去模拟节点
    def walk_by_depth(self, depth):
        pass

    # 回溯
    def backward(self, node: Node, num: int, value: float) -> None:
        node.value += value
        node.times += num
        # 不是父节点就继续增加
        if node.parent is not None:
            return self.backward(node.parent, num, value)

    # 计算UCB值 (value / n) + sqrt(ln总次数 / 当前节点探索次数)
    # 实际不可能出现复数 因为all_num永远会大于0-1
    @staticmethod
    def calc_UCB(node: Node, all_num: int, color: evaluate.Color) -> float:
        value: float = node.value
        n: int = node.times

        # 没访问过 ucb值是无穷大
        if n == 0 or all_num == 0:
            return float('inf')

        # 如果是复数则只保留实数部分
        if color == evaluate.Color.WHITE:
            res = value / n + cmath.sqrt(math.log(all_num) / n)
        else:
            res = value / n - cmath.sqrt(math.log(all_num) / n)
        return res.real

    # 判断是否是叶子节点
    @staticmethod
    def is_leaf(node: Node) -> bool:
        if len(node.leaf) == 0:
            return True
        else:
            return False
